/**
 * AI助手查询处理系统
 * 主要功能：
 * 1. 处理与Claude AI的对话交互
 * 2. 管理工具使用（Tool Use）流程
 * 3. 处理二进制反馈机制
 * 4. 权限检查和工具执行
 */

import {
  Message as APIAssistantMessage,
  MessageParam,
  ToolUseBlock,
} from '@anthropic-ai/sdk/resources/index.mjs'
import { UUID } from 'crypto'
import type { Tool, ToolUseContext } from './Tool.js'
import {
  messagePairValidForBinaryFeedback,
  shouldUseBinaryFeedback,
} from './components/binary-feedback/utils.js'
import { CanUseToolFn } from './hooks/useCanUseTool.js'
import {
  formatSystemPromptWithContext,
  querySonnet,
} from './services/claude.js'
import { logEvent } from './services/statsig.js'
import { all } from './utils/generators.js'
import { logError } from './utils/log.js'
import {
  createAssistantMessage,
  createProgressMessage,
  createToolResultStopMessage,
  createUserMessage,
  FullToolUseResult,
  INTERRUPT_MESSAGE,
  INTERRUPT_MESSAGE_FOR_TOOL_USE,
  NormalizedMessage,
  normalizeMessagesForAPI,
} from './utils/messages.js'
import { BashTool } from './tools/BashTool/BashTool.js'
import { getCwd } from './utils/state.js'

// 响应类型定义
export type Response = { costUSD: number; response: string }

// 用户消息类型定义
export type UserMessage = {
  message: MessageParam
  type: 'user'
  uuid: UUID
  toolUseResult?: FullToolUseResult
}

// 助手消息类型定义
export type AssistantMessage = {
  costUSD: number
  durationMs: number
  message: APIAssistantMessage
  type: 'assistant'
  uuid: UUID
  isApiErrorMessage?: boolean
}

// 二进制反馈结果类型定义
export type BinaryFeedbackResult =
  | { message: AssistantMessage | null; shouldSkipPermissionCheck: false }
  | { message: AssistantMessage; shouldSkipPermissionCheck: true }

// 进度消息类型定义
export type ProgressMessage = {
  content: AssistantMessage
  normalizedMessages: NormalizedMessage[]
  siblingToolUseIDs: Set<string>
  tools: Tool[]
  toolUseID: string
  type: 'progress'
  uuid: UUID
}

// 消息联合类型 - 每个数组项要么是单个消息，要么是消息-响应对
export type Message = UserMessage | AssistantMessage | ProgressMessage

// 最大并发工具使用数量
const MAX_TOOL_USE_CONCURRENCY = 10

/**
 * 带二进制反馈的查询处理函数
 * 功能：处理AI助手的响应，支持二进制反馈机制来改进响应质量
 * 
 * @param toolUseContext - 工具使用上下文
 * @param getAssistantResponse - 获取助手响应的函数
 * @param getBinaryFeedbackResponse - 可选的二进制反馈响应函数
 * @returns 返回一个消息或null（如果用户取消）
 */
async function queryWithBinaryFeedback(
  toolUseContext: ToolUseContext,
  getAssistantResponse: () => Promise<AssistantMessage>,
  getBinaryFeedbackResponse?: (
    m1: AssistantMessage,
    m2: AssistantMessage,
  ) => Promise<BinaryFeedbackResult>,
): Promise<BinaryFeedbackResult> {
  // 检查是否应该使用二进制反馈机制
  if (
    process.env.USER_TYPE !== 'ant' ||
    !getBinaryFeedbackResponse ||
    !(await shouldUseBinaryFeedback())
  ) {
    // 不使用二进制反馈，直接获取单个响应
    const assistantMessage = await getAssistantResponse()
    if (toolUseContext.abortController.signal.aborted) {
      return { message: null, shouldSkipPermissionCheck: false }
    }
    return { message: assistantMessage, shouldSkipPermissionCheck: false }
  }
  
  // 使用二进制反馈：并行获取两个响应
  const [m1, m2] = await Promise.all([
    getAssistantResponse(),
    getAssistantResponse(),
  ])
  
  if (toolUseContext.abortController.signal.aborted) {
    return { message: null, shouldSkipPermissionCheck: false }
  }
  
  // 错误处理逻辑
  if (m2.isApiErrorMessage) {
    // 如果m2是错误，返回m1（即使m1也是错误）
    return { message: m1, shouldSkipPermissionCheck: false }
  }
  if (m1.isApiErrorMessage) {
    return { message: m2, shouldSkipPermissionCheck: false }
  }
  
  // 检查消息对是否适合二进制反馈
  if (!messagePairValidForBinaryFeedback(m1, m2)) {
    return { message: m1, shouldSkipPermissionCheck: false }
  }
  
  // 使用二进制反馈选择更好的响应
  return await getBinaryFeedbackResponse(m1, m2)
}

/**
 * 思考规则说明（来自注释的翻译）：
 * 思考规则是冗长而幸运的。它们需要大量的思考和长时间的深度冥想，
 * 让一个巫师能够理解这些规则。
 * 
 * 规则如下：
 * 1. 包含thinking或redacted_thinking块的消息必须是max_thinking_length > 0的查询的一部分
 * 2. thinking块不能是块中的最后一条消息
 * 3. thinking块必须在助手轨迹的持续时间内保持（单轮对话，或如果该轮包含tool_use块，
 *    则还包括其后续的tool_result和以下助手消息）
 * 
 * 好好遵守这些规则，年轻的巫师。因为它们是思考的规则，
 * 而思考的规则就是宇宙的规则。如果你不遵守这些规则，
 * 你将受到一整天的调试和抓头发的惩罚。
 */

/**
 * 主要的查询处理函数
 * 功能：处理AI对话流程，包括消息处理、工具使用和递归查询
 * 
 * @param messages - 消息历史
 * @param systemPrompt - 系统提示词
 * @param context - 上下文信息
 * @param canUseTool - 工具使用权限检查函数
 * @param toolUseContext - 工具使用上下文
 * @param getBinaryFeedbackResponse - 可选的二进制反馈函数
 * @returns 异步生成器，产生消息流
 */
export async function* query(
  messages: Message[],
  systemPrompt: string[],
  context: { [k: string]: string },
  canUseTool: CanUseToolFn,
  toolUseContext: ToolUseContext,
  getBinaryFeedbackResponse?: (
    m1: AssistantMessage,
    m2: AssistantMessage,
  ) => Promise<BinaryFeedbackResult>,
): AsyncGenerator<Message, void> {
  // 格式化完整的系统提示词
  const fullSystemPrompt = formatSystemPromptWithContext(systemPrompt, context)

  /**
   * 获取助手响应的内部函数
   * 调用Claude API获取AI助手的响应
   */
  function getAssistantResponse() {
    return querySonnet(
      normalizeMessagesForAPI(messages),
      fullSystemPrompt,
      toolUseContext.options.maxThinkingTokens,
      toolUseContext.options.tools,
      toolUseContext.abortController.signal,
      {
        dangerouslySkipPermissions:
          toolUseContext.options.dangerouslySkipPermissions ?? false,
        model: toolUseContext.options.slowAndCapableModel,
        prependCLISysprompt: true,
      },
    )
  }

  // 获取助手响应（可能使用二进制反馈）
  const result = await queryWithBinaryFeedback(
    toolUseContext,
    getAssistantResponse,
    getBinaryFeedbackResponse,
  )

  // 检查是否被中断
  if (result.message === null) {
    yield createAssistantMessage(INTERRUPT_MESSAGE)
    return
  }

  const assistantMessage = result.message
  const shouldSkipPermissionCheck = result.shouldSkipPermissionCheck

  // 输出助手消息
  yield assistantMessage

  // 提取工具使用消息
  // 注意：stop_reason === 'tool_use'不可靠，并不总是正确设置
  const toolUseMessages = assistantMessage.message.content.filter(
    _ => _.type === 'tool_use',
  )

  // 如果没有工具使用，查询结束
  if (!toolUseMessages.length) {
    return
  }

  const toolResults: UserMessage[] = []

  // 优先并发运行工具（如果可以的话）
  // TODO: 收紧逻辑 - 我们可以比这更频繁地并发运行
  if (
    toolUseMessages.every(msg =>
      toolUseContext.options.tools.find(t => t.name === msg.name)?.isReadOnly(),
    )
  ) {
    // 所有工具都是只读的，可以并发执行
    for await (const message of runToolsConcurrently(
      toolUseMessages,
      assistantMessage,
      canUseTool,
      toolUseContext,
      shouldSkipPermissionCheck,
    )) {
      yield message
      // 进度消息不会发送到服务器，所以不需要为下一轮累积
      if (message.type === 'user') {
        toolResults.push(message)
      }
    }
  } else {
    // 有非只读工具，需要串行执行
    for await (const message of runToolsSerially(
      toolUseMessages,
      assistantMessage,
      canUseTool,
      toolUseContext,
      shouldSkipPermissionCheck,
    )) {
      yield message
      // 进度消息不会发送到服务器，所以不需要为下一轮累积
      if (message.type === 'user') {
        toolResults.push(message)
      }
    }
  }

  // 检查是否在工具使用过程中被中断
  if (toolUseContext.abortController.signal.aborted) {
    yield createAssistantMessage(INTERRUPT_MESSAGE_FOR_TOOL_USE)
    return
  }

  // 对工具结果进行排序，以匹配toolUseMessages的顺序
  const orderedToolResults = toolResults.sort((a, b) => {
    const aIndex = toolUseMessages.findIndex(
      tu => tu.id === (a.message.content[0] as ToolUseBlock).id,
    )
    const bIndex = toolUseMessages.findIndex(
      tu => tu.id === (b.message.content[0] as ToolUseBlock).id,
    )
    return aIndex - bIndex
  })

  // 递归调用query函数，继续处理对话
  yield* await query(
    [...messages, assistantMessage, ...orderedToolResults],
    systemPrompt,
    context,
    canUseTool,
    toolUseContext,
    getBinaryFeedbackResponse,
  )
}

/**
 * 并发运行工具
 * 功能：同时执行多个工具，提高效率
 * 
 * @param toolUseMessages - 工具使用消息数组
 * @param assistantMessage - 助手消息
 * @param canUseTool - 工具使用权限检查函数
 * @param toolUseContext - 工具使用上下文
 * @param shouldSkipPermissionCheck - 是否跳过权限检查
 * @returns 异步生成器，产生消息流
 */
async function* runToolsConcurrently(
  toolUseMessages: ToolUseBlock[],
  assistantMessage: AssistantMessage,
  canUseTool: CanUseToolFn,
  toolUseContext: ToolUseContext,
  shouldSkipPermissionCheck?: boolean,
): AsyncGenerator<Message, void> {
  yield* all(
    toolUseMessages.map(toolUse =>
      runToolUse(
        toolUse,
        new Set(toolUseMessages.map(_ => _.id)),
        assistantMessage,
        canUseTool,
        toolUseContext,
        shouldSkipPermissionCheck,
      ),
    ),
    MAX_TOOL_USE_CONCURRENCY,
  )
}

/**
 * 串行运行工具
 * 功能：按顺序执行工具，确保安全性
 * 
 * @param toolUseMessages - 工具使用消息数组
 * @param assistantMessage - 助手消息
 * @param canUseTool - 工具使用权限检查函数
 * @param toolUseContext - 工具使用上下文
 * @param shouldSkipPermissionCheck - 是否跳过权限检查
 * @returns 异步生成器，产生消息流
 */
async function* runToolsSerially(
  toolUseMessages: ToolUseBlock[],
  assistantMessage: AssistantMessage,
  canUseTool: CanUseToolFn,
  toolUseContext: ToolUseContext,
  shouldSkipPermissionCheck?: boolean,
): AsyncGenerator<Message, void> {
  for (const toolUse of toolUseMessages) {
    yield* runToolUse(
      toolUse,
      new Set(toolUseMessages.map(_ => _.id)),
      assistantMessage,
      canUseTool,
      toolUseContext,
      shouldSkipPermissionCheck,
    )
  }
}

/**
 * 运行单个工具使用
 * 功能：执行单个工具，包括错误处理和权限检查
 * 
 * @param toolUse - 工具使用块
 * @param siblingToolUseIDs - 兄弟工具使用ID集合
 * @param assistantMessage - 助手消息
 * @param canUseTool - 工具使用权限检查函数
 * @param toolUseContext - 工具使用上下文
 * @param shouldSkipPermissionCheck - 是否跳过权限检查
 * @returns 异步生成器，产生消息流
 */
export async function* runToolUse(
  toolUse: ToolUseBlock,
  siblingToolUseIDs: Set<string>,
  assistantMessage: AssistantMessage,
  canUseTool: CanUseToolFn,
  toolUseContext: ToolUseContext,
  shouldSkipPermissionCheck?: boolean,
): AsyncGenerator<Message, void> {
  const toolName = toolUse.name
  const tool = toolUseContext.options.tools.find(t => t.name === toolName)

  // 检查工具是否存在
  if (!tool) {
    logEvent('tengu_tool_use_error', {
      error: `No such tool available: ${toolName}`,
      messageID: assistantMessage.message.id,
      toolName,
      toolUseID: toolUse.id,
    })
    yield createUserMessage([
      {
        type: 'tool_result',
        content: `Error: No such tool available: ${toolName}`,
        is_error: true,
        tool_use_id: toolUse.id,
      },
    ])
    return
  }

  const toolInput = toolUse.input as { [key: string]: string }

  try {
    // 检查是否被中断
    if (toolUseContext.abortController.signal.aborted) {
      logEvent('tengu_tool_use_cancelled', {
        toolName: tool.name,
        toolUseID: toolUse.id,
      })
      const message = createUserMessage([
        createToolResultStopMessage(toolUse.id),
      ])
      yield message
      return
    }

    // 执行权限检查和工具调用
    for await (const message of checkPermissionsAndCallTool(
      tool,
      toolUse.id,
      siblingToolUseIDs,
      toolInput,
      toolUseContext,
      canUseTool,
      assistantMessage,
      shouldSkipPermissionCheck,
    )) {
      yield message
    }
  } catch (e) {
    logError(e)
  }
}

/**
 * 工具输入标准化
 * 功能：对工具输入进行预处理和标准化
 * TODO: 将此功能推广到所有工具
 * 
 * @param tool - 工具对象
 * @param input - 输入参数
 * @returns 标准化后的输入
 */
export function normalizeToolInput(
  tool: Tool,
  input: { [key: string]: boolean | string | number },
): { [key: string]: boolean | string | number } {
  switch (tool) {
    case BashTool: {
      const { command, timeout } = BashTool.inputSchema.parse(input) // 已在上游验证，不会抛出异常
      return {
        command: command.replace(`cd ${getCwd()} && `, ''),
        ...(timeout ? { timeout } : {}),
      }
    }
    default:
      return input
  }
}

/**
 * 权限检查和工具调用
 * 功能：验证输入、检查权限并执行工具
 * 
 * @param tool - 工具对象
 * @param toolUseID - 工具使用ID
 * @param siblingToolUseIDs - 兄弟工具使用ID集合
 * @param input - 输入参数
 * @param context - 工具使用上下文
 * @param canUseTool - 工具使用权限检查函数
 * @param assistantMessage - 助手消息
 * @param shouldSkipPermissionCheck - 是否跳过权限检查
 * @returns 异步生成器，产生用户消息或进度消息
 */
async function* checkPermissionsAndCallTool(
  tool: Tool,
  toolUseID: string,
  siblingToolUseIDs: Set<string>,
  input: { [key: string]: boolean | string | number },
  context: ToolUseContext,
  canUseTool: CanUseToolFn,
  assistantMessage: AssistantMessage,
  shouldSkipPermissionCheck?: boolean,
): AsyncGenerator<UserMessage | ProgressMessage, void> {
  // 使用zod验证输入类型
  // （令人惊讶的是，模型在生成有效输入方面并不擅长）
  const isValidInput = tool.inputSchema.safeParse(input)
  if (!isValidInput.success) {
    logEvent('tengu_tool_use_error', {
      error: `InputValidationError: ${isValidInput.error.message}`,
      messageID: assistantMessage.message.id,
      toolName: tool.name,
      toolInput: JSON.stringify(input).slice(0, 200),
    })
    yield createUserMessage([
      {
        type: 'tool_result',
        content: `InputValidationError: ${isValidInput.error.message}`,
        is_error: true,
        tool_use_id: toolUseID,
      },
    ])
    return
  }

  // 标准化工具输入
  const normalizedInput = normalizeToolInput(tool, input)

  // 验证输入值。每个工具都有自己的验证逻辑
  const isValidCall = await tool.validateInput?.(
    normalizedInput as never,
    context,
  )
  if (isValidCall?.result === false) {
    logEvent('tengu_tool_use_error', {
      error: isValidCall?.message.slice(0, 2000),
      messageID: assistantMessage.message.id,
      toolName: tool.name,
      toolInput: JSON.stringify(input).slice(0, 200),
      ...(isValidCall?.meta ?? {}),
    })
    yield createUserMessage([
      {
        type: 'tool_result',
        content: isValidCall!.message,
        is_error: true,
        tool_use_id: toolUseID,
      },
    ])
    return
  }

  // 检查是否有权限使用工具，如果没有则向用户请求权限
  const permissionResult = shouldSkipPermissionCheck
    ? ({ result: true } as const)
    : await canUseTool(tool, normalizedInput, context, assistantMessage)
  if (permissionResult.result === false) {
    yield createUserMessage([
      {
        type: 'tool_result',
        content: permissionResult.message,
        is_error: true,
        tool_use_id: toolUseID,
      },
    ])
    return
  }

  // 调用工具
  try {
    const generator = tool.call(normalizedInput as never, context, canUseTool)
    for await (const result of generator) {
      switch (result.type) {
        case 'result':
          // 工具执行成功
          logEvent('tengu_tool_use_success', {
            messageID: assistantMessage.message.id,
            toolName: tool.name,
          })
          yield createUserMessage(
            [
              {
                type: 'tool_result',
                content: result.resultForAssistant,
                tool_use_id: toolUseID,
              },
            ],
            {
              data: result.data,
              resultForAssistant: result.resultForAssistant,
            },
          )
          return
        case 'progress':
          // 工具执行进度更新
          logEvent('tengu_tool_use_progress', {
            messageID: assistantMessage.message.id,
            toolName: tool.name,
          })
          yield createProgressMessage(
            toolUseID,
            siblingToolUseIDs,
            result.content,
            result.normalizedMessages,
            result.tools,
          )
      }
    }
  } catch (error) {
    // 工具执行出错
    const content = formatError(error)
    logError(error)
    logEvent('tengu_tool_use_error', {
      error: content.slice(0, 2000),
      messageID: assistantMessage.message.id,
      toolName: tool.name,
      toolInput: JSON.stringify(input).slice(0, 1000),
    })
    yield createUserMessage([
      {
        type: 'tool_result',
        content,
        is_error: true,
        tool_use_id: toolUseID,
      },
    ])
  }
}

/**
 * 错误格式化函数
 * 功能：将错误对象格式化为可读的字符串，包括stdout和stderr
 * 
 * @param error - 错误对象
 * @returns 格式化后的错误字符串
 */
function formatError(error: unknown): string {
  if (!(error instanceof Error)) {
    return String(error)
  }
  const parts = [error.message]
  if ('stderr' in error && typeof error.stderr === 'string') {
    parts.push(error.stderr)
  }
  if ('stdout' in error && typeof error.stdout === 'string') {
    parts.push(error.stdout)
  }
  const fullMessage = parts.filter(Boolean).join('\n')
  if (fullMessage.length <= 10000) {
    return fullMessage
  }
  // 如果消息太长，进行截断处理
  const halfLength = 5000
  const start = fullMessage.slice(0, halfLength)
  const end = fullMessage.slice(-halfLength)
  return `${start}\n\n... [${fullMessage.length - 10000} characters truncated] ...\n\n${end}`
}
